﻿using KaldiNativeFbankSharp;
using TorchSharp;

namespace SenseVoiceSharp
{
    public class WavFrontend
    {
        OnlineFbank _onlineFbank;
        torch.Tensor _cmvn;
        int _lfrM = 7;
        int _lfrN = 6;

        public WavFrontend(string cmvnFile, int lfrM = 7, int lftN = 6)
        {
            _onlineFbank = new OnlineFbank(0, true, 16000, 80, 10, 25, 0, false);
            _cmvn = LoadCmvn(cmvnFile);
            _lfrM = lfrM;
            _lfrN = lftN;
        }

        public torch.Tensor GetFeatures(torch.Tensor inputs)
        {
            var (fbank, fbank_len) = Fbank(inputs); 
            var feats = ApplyCmvn(ApplyLfr(fbank, _lfrM, _lfrN));
            return feats;
        }

        public (torch.Tensor, long) Fbank(torch.Tensor waveform)
        {
            waveform = waveform * 32768;
            float[] waveformlist = waveform.data<float>().ToArray();
            var fbank = _onlineFbank.GetFbank(waveformlist);
            var feat = torch.tensor(fbank).reshape(-1, 80);
            var feat_len = feat.shape[0];
            return (feat, feat_len);
        }

        public (torch.Tensor, long) Lfr_cmvn(torch.Tensor feat)
        {
            if (_lfrM != 1 || _lfrN != 1)
            {
                feat = ApplyLfr(feat, _lfrM, _lfrN);
            }
            feat = ApplyCmvn(feat);
            long feat_len = feat.shape[0];
            return (feat, feat_len);
        }

        public torch.Tensor ApplyCmvn(torch.Tensor inputs)
        {
            var frame = inputs.shape[0];
            var dim = inputs.shape[1];
            var means = torch.tile(_cmvn[torch.TensorIndex.Slice(0, 1), torch.TensorIndex.Slice(null, dim)], new long[] { frame, 1 });
            var vars = torch.tile(_cmvn[torch.TensorIndex.Slice(1, 2), torch.TensorIndex.Slice(null, dim)], new long[] { frame, 1 });
            inputs = (inputs + means) * vars;
            return inputs;
        }

        public torch.Tensor ApplyLfr(torch.Tensor inputs, int lfrM, int lfrN)
        {
            var T = inputs.shape[0];
            var T_lfr = (int)Math.Ceiling((float)T / lfrN);
            var leftPadding = torch.tile(inputs[0], new long[] { (lfrM - 1) / 2, 1 });
            inputs = torch.vstack(new List<torch.Tensor> { leftPadding, inputs });
            T = T + (lfrM - 1) / 2;
            List<torch.Tensor> lfrLists = new List<torch.Tensor>();
            
            for (int i = 0; i < T_lfr; i++)
            {
                if (lfrM < T - i * lfrN)
                {
                    var index = i * lfrN;
                    var index1 = i * lfrN + lfrM;
                    lfrLists.Add(inputs[torch.TensorIndex.Slice(index, index1)].reshape(1, -1));
                }
                else
                {
                    var numPadding = lfrM - (T - i * lfrN);
                    var index = i * lfrN;
                    var frame = inputs[torch.TensorIndex.Slice(index, null)].reshape(-1);                    
                    for (int j = 0; j < numPadding; j++)
                    {
                        frame = torch.hstack(new List<torch.Tensor> { frame, inputs[-1] });
                    }
                   
                    lfrLists.Add(frame);
                }
            }
            var ret = torch.vstack(lfrLists).to(torch.float32);
            return ret;
        }

        public torch.Tensor LoadCmvn(string cmvnFile)
        {
            var tool = new CMVNFileParser(cmvnFile);
            var meanslists = new List<string>();
            var varslists = new List<string>();
            tool.Parse(out meanslists, out varslists);
            var meanslistsFloat = meanslists.Select(x => float.Parse(x)).ToArray();
            var varslistssFloat = varslists.Select(x => float.Parse(x)).ToArray();
            var means = torch.tensor(meanslistsFloat);
            var vars = torch.tensor(varslistssFloat);
            var cmvn = torch.vstack(new List<torch.Tensor> { means, vars });
            return cmvn;
        }
    }
}
